"""Data collection and management module for Tianyan quantum platform.

This module handles experiment data collection, submission and result processing
for the Tianyan quantum computing platform.
"""

from datetime import datetime
from typing import List, Optional

import pandas as pd
from cqlib import TianYanPlatform
from cqlib.circuits import Circuit

from .tianyan_platform import (
    send_job, 
    get_result_less_15qubits, 
    get_result_more_15qubits
)


class TianyanData:
    """Handler for Tianyan platform experiment data.
    
    This class manages the submission, collection and storage of quantum experiment
    data from the Tianyan platform.
    
    Attributes:
        num_qubits: Number of qubits in experiments
        num_shots: Number of shots per experiment
        platform: TianYan platform instance
        exp_name: Experiment name with timestamp
        df: DataFrame storing experiment results
    """

    def __init__(self, 
                 num_qubits: int,
                 num_shots: int,
                 platform: TianYanPlatform,
                 exp_name: str):
        """Initialize the data handler.

        Args:
            num_qubits: Number of qubits in experiments
            num_shots: Number of shots per experiment
            platform: TianYan platform instance
            exp_name: Base name for experiments
        """
        self.num_qubits = num_qubits
        self.num_shots = num_shots
        self.platform = platform
        self._exp_name = f"{exp_name}"
        
        # Initialize results dataframe
        self.df = pd.DataFrame(columns=[
            'lab_id', 'exp_id', 'exp_name', 'qubits', 'cbits',
            'true_value', 'exp_result'
        ])
    
    def submit_experiment(self, circuit_list: List[Circuit]) -> None:
        """Submit quantum circuits for execution.

        Args:
            circuit_list: List of quantum circuits to execute
        """
        # Generate unique experiment name with timestamp and microseconds for fast submissions
        self.exp_name = f"{self._exp_name}_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"
        
        # Create new lab for experiments with unique name
        lab_id = self.platform.create_lab(
            name=self.exp_name, 
            remark=self.exp_name
        )

        # Submit each circuit
        for idx, circuit in enumerate(circuit_list):
            exp_id = send_job(
                circuit=circuit,
                lab_id=lab_id,
                nshots=self.num_shots,
                platform=self.platform
            )
            
            # Record experiment metadata
            self.df.loc[idx, 'exp_id'] = exp_id
            # self.df.loc[idx, 'lab_id'] = lab_id
            self.df.loc[idx, 'exp_name'] = self.exp_name

        print('Job submission completed.')

    def receive_result(self) -> None:
        """Collect experiment results from platform."""
        if self.num_qubits < 15:
            get_result_less_15qubits(self.df, self.platform)
        else:
            get_result_more_15qubits(self.df, self.platform)
    
    def save_result(self, file_name: str) -> None:
        """Save results to CSV file.

        Args:
            file_name: Path to output CSV file
        """
        self.df.to_csv(file_name, index=False)
    